import queue
import threading
import time
import random

class MultiThreadPool:
    def __init__(self, thread_num, msg_list, worker_function):
        """
        初始化线程池
        :param thread_num: 线程数量
        :param msg_list: 待处理的消息列表
        :param worker_function: 处理消息的工作函数，接收一个msg参数
        """
        self.thread_num = thread_num
        self.msg_list = msg_list
        self.worker_function = worker_function
        self.queue = queue.Queue()  # queue本身是线程安全的
        self.threads = []
        self._stop_event = threading.Event()  # 使用事件对象更安全地控制线程退出
    
    def _worker(self):
        """工作线程函数，从队列中获取消息并处理"""
        # 循环直到收到事件被设置且队列为空
        while not self._stop_event.is_set() or not self.queue.empty():
            try:
                # 非阻塞超时获取，允许定期检查退出信号
                msg = self.queue.get(timeout=1)
                try:
                    self.worker_function(msg)
                finally:
                    self.queue.task_done()  # 确保即使处理出错也会标记完成
            except queue.Empty:
                continue
    
    def start(self):
        """启动线程池，开始处理消息"""
        # 将所有消息放入队列
        for msg in self.msg_list:
            self.queue.put(msg)
        
        # 重置停止事件
        self._stop_event.clear()
        
        # 创建并启动线程
        for _ in range(self.thread_num):
            t = threading.Thread(target=self._worker)
            t.daemon = True  # 守护线程
            t.start()
            self.threads.append(t)
    
    def wait_end(self):
        """等待所有消息处理完成"""
        # 等待队列中的所有任务完成
        self.queue.join()
        
        # 设置停止事件，通知线程可以退出
        self._stop_event.set()
        
        # 等待所有线程结束
        for t in self.threads:
            t.join()


# 测试函数
def test_worker(msg):
    """测试用的工作函数，模拟处理消息"""
    print(f"线程 {threading.current_thread().name} 开始处理消息: {msg}")


def test_multi_thread_pool():
    """测试多线程池"""
    print("开始测试线程安全的多线程池...")
    
    # 创建测试消息列表
    msg_list = [f"消息 {i}" for i in range(1, 11)]
    
    # 创建线程池，使用3个线程
    thread_pool = MultiThreadPool(
        thread_num=3,
        msg_list=msg_list,
        worker_function=test_worker
    )
    
    # 启动线程池
    start_time = time.time()
    thread_pool.start()
    
    # 等待所有任务完成
    thread_pool.wait_end()
    
    end_time = time.time()
    print(f"所有消息处理完成，总耗时: {end_time - start_time:.2f}秒")


if __name__ == "__main__":
    test_multi_thread_pool()
    